package completion

import (
	"chatgpt/logq"
	"context"
	"github.com/slclub/go-tips"
	"github.com/slclub/go-tips/json"
	"github.com/slclub/go-tips/spinlock"
	"github.com/slclub/go-tips/stringbyte"
	"strings"
	"sync"
	"time"
)

var (
	manager                 *Manager
	DEFAULT_COMPLETION_SIZE = 1024
	DEFAULT_COMPLETION_BOT  = "TD"
	CHAT_USER               = "user"
	CHAT_BOT                = "assistant"
)

type Dialogue struct {
	Prompt     string
	Completion string
	expire     time.Time
}

type DialogueQueue struct {
	Queue  []*Dialogue
	Uid    string
	mu     sync.Locker
	expire time.Time
	Bot    string
	Size   int
}

type Manager struct {
	dialogues  sync.Map
	ctx        context.Context
	EmptyQueue DialogueQueue
}

func Mgr(ctxs ...context.Context) *Manager {
	var ctx context.Context
	if len(ctxs) > 0 {
		ctx = ctxs[0]
	}
	if manager == nil {
		if ctx == nil {
			ctx = context.Background()
		}
		manager = &Manager{
			dialogues: sync.Map{},
			ctx:       ctx,
		}
		manager.Init()
	}
	if ctx != nil {
		manager.ctx = ctx
	}
	return manager
}

// manager
func (self *Manager) Init() {
	self.Watch()
}

func (self *Manager) Watch() {
	go func() {
		t := time.NewTicker(time.Second)
		for {
			select {
			case <-t.C:
				self.Update()
			case <-self.ctx.Done():
				logq.InfoF("COMPLETION.CONVERSATION.Watch quit")
				return
			}
		}
	}()
}

// 定时执行的更新操作
// 是否用锁，待考虑
func (self *Manager) Update() {
	now := time.Now()
	self.dialogues.Range(func(key, value any) bool {
		obj, ok := value.(*DialogueQueue)
		if !ok {
			return true
		}
		if obj.ExpireCheck(now) {
			self.dialogues.Delete(key)
			return true
		}
		obj.Update(now)
		return true
	})
}

func (self *Manager) Add(uid string, dia *Dialogue) {
	obj, ok := self.dialogues.Load(uid)
	if ok {
		q := obj.(*DialogueQueue)
		q.Append(dia)
		return
	}
	q := newDialogueQueue(uid)
	q.Append(dia)
	self.dialogues.Store(uid, q)
}

func (self *Manager) Get(uid string) *DialogueQueue {
	obj, ok := self.dialogues.Load(uid)
	if ok {
		return obj.(*DialogueQueue)
	}

	return nil
}

func (self *Manager) StopSequence(uid string) string {
	queue := self.Get(uid)
	if queue == nil {
		return DEFAULT_COMPLETION_BOT
	}
	return tips.String((queue.Uid)) + "," + queue.Bot
}

// Dialogue Queue
func newDialogueQueue(uid string) *DialogueQueue {
	return &DialogueQueue{
		Queue:  make([]*Dialogue, 0),
		Uid:    uid,
		mu:     spinlock.New(),
		Bot:    DEFAULT_COMPLETION_BOT,
		Size:   DEFAULT_COMPLETION_SIZE,
		expire: time.Now().Add(time.Hour * 3),
	}
}
func (self *DialogueQueue) Len() int {
	return len(self.Queue)
}

func (self *DialogueQueue) Update(now time.Time) {
	n := self.Len()
	if n == 0 {
		return
	}
	self.mu.Lock()
	defer self.mu.Unlock()
	for k, v := range self.Queue {
		if !v.ExpireCheck(now) {
			continue
		}
		if k == n-1 {
			self.Queue = self.Queue[:n-1]
		} else {
			self.Queue = append(self.Queue[:k], self.Queue[k+1:]...)
		}
		n--
	}
}

func (self *DialogueQueue) Append(dia *Dialogue) {
	self.mu.Lock()
	defer self.mu.Unlock()
	self.Queue = append(self.Queue, dia)
	self.expire = time.Now().Add(time.Hour * 3)
}

func (self *DialogueQueue) ExpireCheck(now time.Time) bool {
	if now.Sub(self.expire) > 0 {
		return true
	}
	return false
}

func (self *DialogueQueue) Clear() {
	self.mu.Lock()
	defer self.mu.Unlock()
	self.Queue = []*Dialogue{}
}

func (self *DialogueQueue) ConvPrompt(prompt string) string {
	prompt = "\n" + prompt + " " + self.Bot + ":"
	if self.Len() == 0 {
		return prompt
	}
	uidstr := tips.String((self.Uid))
	data := []string{}
	self.mu.Lock()
	defer self.mu.Unlock()
	size := 0
	for i := self.Len() - 1; i >= 0; i-- {
		v := self.Queue[i]

		// 过长则去掉前面的聊天记录
		size += v.Size()
		if size >= self.Size {
			di := i
			if i < self.Len()-1 {
				di++
			}
			self.Queue = self.Queue[di:]
			break
		}

		// 将正确的聊天记录复加在一起
		data = append(data, v.String(uidstr, self.Bot))
	}
	return strings.Join(data, "\n") + "\n" + uidstr + ":" + prompt
}

func (self *DialogueQueue) RepeatCompletion(prompt string) string {
	if self.Len() == 0 {
		return ""
	}
	trim := func(s string) string {
		ds := strings.Trim(s, " ")
		ds = strings.Trim(s, "\n")
		ds = strings.Trim(s, "\r\n")
		return ds
	}
	data := self.Queue[self.Len()-1]
	if tips.StrPos(prompt, data.Prompt) >= 0 {
		return trim(data.Completion)
	}
	if tips.StrPos(data.Prompt, prompt) >= 0 {
		return trim(data.Completion)
	}
	return ""
}

func (self *DialogueQueue) ConvPromptToChat(prompt string) string {
	//prompt = "\n" + prompt + self.Bot + ":"
	data := []map[string]any{}
	size := 0
	if self.Len() == 0 {
		data = []map[string]any{map[string]any{
			"role":    CHAT_USER,
			"content": prompt,
		}}
		goto TAIL
	}

	self.mu.Lock()
	defer self.mu.Unlock()

	for i := self.Len() - 1; i >= 0; i-- {
		v := self.Queue[i]

		// 过长则去掉前面的聊天记录
		size += v.Size()
		if size >= self.Size {
			di := i
			if i < self.Len()-1 {
				di++
			}
			self.Queue = self.Queue[di:]
			break
		}

		// 将正确的聊天记录复加在一起
		data = append(data, map[string]any{
			"role":    CHAT_USER,
			"content": v.Prompt,
		}, map[string]any{
			"role":    CHAT_BOT,
			"content": v.Completion,
		})
	}
	data = append(data, map[string]any{
		"role":    CHAT_USER,
		"content": prompt,
	})
TAIL:
	datab, err := json.Marshal(data)
	if err != nil {
		return ""
	}
	return stringbyte.BytesToString(datab)
}

// Dialogue
func NewDialogue(prompt, completion string) *Dialogue {
	return &Dialogue{
		Prompt:     prompt,
		Completion: completion,
		expire:     time.Now().Add(time.Hour * 2),
	}
}
func (self *Dialogue) ExpireCheck(now time.Time) bool {
	if now.Sub(self.expire) > 0 {
		return true
	}
	return false
}

func (self *Dialogue) String(uid, bot string) string {
	return uid + ":" + self.Prompt + "\n" + bot + ":" + self.Completion
}

func (self *Dialogue) Size() int {
	size := len(self.Prompt) + len(self.Completion)
	return size * 2
}
